// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/vector/complex_vector.h>
#include <pollux/vector/vector_maker.h>

#include <melon/benchmark.h>
#include <melon/init/init.h>

#include <random>

TURBO_FLAG(int32_t,seed, 0, "Random seed");

namespace kumo::pollux {
namespace {

MapVectorPtr generateBase(
    memory::MemoryPool* pool,
    std::default_random_engine& gen) {
  constexpr int kSize = 1000;
  constexpr int kMapSize = 10;
  VectorMaker maker(pool);
  std::uniform_real_distribution valueDist;
  return maker.map_vector<int64_t, double>(
      kSize,
      [&](auto /*i*/) { return kMapSize; },
      [](auto /*i*/, auto j) { return j; },
      [&](auto /*i*/, auto /*j*/) { return valueDist(gen); });
}

MapVectorPtr generateUpdate(
    memory::MemoryPool* pool,
    std::default_random_engine& gen,
    const MapVector& base) {
  constexpr int kMapSize = 5;
  auto offsets = allocate_indices(base.size(), pool);
  auto* rawOffsets = offsets->asMutable<vector_size_t>();
  rawOffsets[0] = 0;
  for (int i = 1; i < base.size(); ++i) {
    rawOffsets[i] = rawOffsets[i - 1] + kMapSize;
  }
  auto sizes = allocate_indices(base.size(), pool);
  auto* rawSizes = sizes->asMutable<vector_size_t>();
  std::fill(rawSizes, rawSizes + base.size(), kMapSize);
  int64_t keyCandidates[10];
  std::iota(std::begin(keyCandidates), std::end(keyCandidates), 0);
  auto keys = BaseVector::create<FlatVector<int64_t>>(
      BIGINT(), kMapSize * base.size(), pool);
  for (int i = 0; i < base.size(); ++i) {
    std::shuffle(std::begin(keyCandidates), std::end(keyCandidates), gen);
    for (int j = 0; j < kMapSize; ++j) {
      keys->set(j + i * kMapSize, keyCandidates[j]);
    }
  }
  VectorMaker maker(pool);
  std::uniform_real_distribution valueDist;
  auto values = maker.flat_vector<double>(
      kMapSize * base.size(), [&](auto /*i*/) { return valueDist(gen); });
  return std::make_shared<MapVector>(
      pool,
      base.type(),
      nullptr,
      base.size(),
      std::move(offsets),
      std::move(sizes),
      std::move(keys),
      std::move(values));
}

} // namespace
} // namespace kumo::pollux

int main(int argc, char* argv[]) {
  using namespace kumo::pollux;
  melon::Init melonInit(&argc, &argv);
  if (gflags::GetCommandLineFlagInfoOrDie("seed").is_default) {
    FLAGS_seed = std::random_device{}();
    KLOG(INFO) << "Use generated random seed " << FLAGS_seed;
  }

  memory::MemoryManager::initialize({});
  auto pool = memory::memoryManager()->addLeafPool();
  std::default_random_engine gen(FLAGS_seed);
  auto base = generateBase(pool.get(), gen);
  auto update = generateUpdate(pool.get(), gen, *base);

  melon::addBenchmark(__FILE__, "int64 keys", [&] {
    auto result = base->update({update});
    melon::doNotOptimizeAway(result);
    return base->size();
  });
  melon::runBenchmarks();

  return 0;
}
